from keras import backend as K
import numpy as np
import tensorflow as tf
from PIL import Image

def get_edge_img(filename, image_width = 256, image_height = 256, target_regression = True):
    edgemaps = []

    # 边缘图
    em = Image.open(filename)

    # 调整为指定尺寸大小
    em = em.resize((image_width, image_height))

    em = np.array(em.convert('L'), dtype=np.float32)

    if target_regression:
        bin_em = em / 255.0
    else:
        bin_em = np.zeros_like(em)
        bin_em[np.where(em)] = 1

    # Some edge maps have 3 channels some dont
    bin_em = bin_em if bin_em.ndim == 2 else bin_em[:, :, 0]
    # To fit [batch_size, H, W, 1] output of the network
    bin_em = np.expand_dims(bin_em, 2)

    edgemaps.append(bin_em)
    edgemaps = np.asarray(edgemaps)

    return edgemaps

def _to_tensor(x, dtype):
    """Convert the input `x` to a tensor of type `dtype`.
    # Arguments
    x: An object to be converted (numpy array, list, tensors).
    dtype: The destination type.
    # Returns
    A tensor.
    """
    x = tf.convert_to_tensor(x)  # 转为Tensor向量
    if x.dtype != dtype:
        x = tf.cast(x, dtype)  # 如果和pred向量不是一种数据类型，则转化
    return x

def cross_entropy_balanced(y_true, y_pred):
    _epsilon = _to_tensor(K.epsilon(), y_pred.dtype.base_dtype)
    y_pred = tf.clip_by_value(y_pred, _epsilon, 1 - _epsilon)  # 把张量的每个元素值都压缩在min和max之间，这里相当于0,1之间
    y_pred = tf.log(y_pred / (1 - y_pred))  # 求 y_pred / (1 - y_pred) 的对数

    y_true = tf.cast(y_true, tf.float32)  # 转化y_true为张量形式

    # 对所有维度的值求和！
    count_neg = tf.reduce_sum(1. - y_true)   # 求出负样本（非边缘）有多少点
    count_pos = tf.reduce_sum(y_true)   # 求出正样本（边缘）有多少点

    beta = count_neg / (count_neg + count_pos)  # 因为边缘相对于非边缘很少，所以各自加上相应的权重来保证损失值平衡

    pos_weight = beta / (1 - beta)  # 求一个权重

    # 计算具有权重的sigmoid交叉熵函数
    cost = tf.nn.weighted_cross_entropy_with_logits(logits=y_pred, targets=y_true, pos_weight=pos_weight)

    # Multiply by 1 - beta
    cost = tf.reduce_mean(cost * (1 - beta))   # 对所有维度的值求平均值！

    return tf.where(tf.equal(count_pos, 0.0), 0.0, cost)

# 像素差，越小越好
def ofuse_pixel_error(y_true, y_pred):
    pred = tf.cast(tf.greater(y_pred, 0.5), tf.int32, name='predictions') # 应该是每个点与0.5比较，返回每个点的布尔值，再转化为int32类型的张量
    error = tf.cast(tf.not_equal(pred, tf.cast(y_true, tf.int32)), tf.float32)  # 应该是和原图，得出不同的点
    return tf.reduce_mean(error, name='pixel_error')  # 对全局求均值，也即像素差

image_path = r'F:\临时存放地\标准数据集\HED-BSDS\train\aug_gt_scale_0.5\45.0_1_0\28096.png'
y_true = get_edge_img(image_path)

y_pred = tf.placeholder(tf.float32, shape=[1, 256,256,1], name='x')

# cross_entropy_balanced(y_true, y_pred)

ofuse_pixel_error(y_true, y_pred)

